''' 用pytorch读取tensorflow flower数据集（2021.8.25 王耀）'''

import torch
import  os
import torchvision
from torchvision import transforms, utils
import  random
from shutil import copy

def mkfile(file):
    if not os.path.exists(file):
        os.makedirs(file)

def split_data(file):
    """
    对tf flowers数据集进行测试集和验证集的分类
    输入：file tfflowers的根目录
    输出：None
    """
    if os.path.exists('./flower_photos/train/') and os.path.exists('./flower_photos/val/'):
        print('return...')
        return

    flower_class = [cla for cla in os.listdir(file) if (".txt" not in cla) and ('.DS' not in cla)]
    mkfile('./flower_photos/train')
    for cla in flower_class:
        mkfile('./flower_photos/train/' + cla)

    mkfile('./flower_photos/val')
    for cla in flower_class:
        mkfile('./flower_photos/val/' + cla)

    split_rate = 0.1
    for cla in flower_class:
        cla_path = file + '/' + cla + '/'
        images = os.listdir(cla_path)
        num = len(images)
        eval_index = random.sample(images, k = int(num * split_rate))
        for index ,ima in enumerate(images):
            image_path = cla_path + ima
            if ima in eval_index:
                new_path = './flower_photos/val/' + cla
                copy(image_path, new_path)
            else:
                new_path = './flower_photos/train/' + cla
                copy(image_path, new_path)

    print('split datasets successfully')



def tf_flower_dl(types, params):
    """
    输入：
    types: <string> "dev" or "train"
    “dev”表示返回测试集，“train”表示返回训练集

    params: 参数
    """
    img_data_train = torchvision.datasets.ImageFolder(
        './flower_photos/train',
        transform=transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor()
        ])
    )

    img_data_val = torchvision.datasets.ImageFolder(
        './flower_photos/val',
        transform=transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor()
        ])
    )

    trainloader = torch.utils.data.DataLoader(img_data_train, batch_size=params.batch_size, shuffle=True,
                                              num_workers=params.num_workers, pin_memory=params.cuda)

    valloader = torch.utils.data.DataLoader(img_data_val, batch_size=params.batch_size, shuffle=True,
                                              num_workers=params.num_workers, pin_memory=params.cuda)

    if types == 'train':
        dl = trainloader
    else:
        dl = valloader

    return dl